package middlewares

import (
	"fmt"
	"gin-study/app/http/responses"
	"gin-study/pkg/helper"
	"github.com/gin-gonic/gin"
	"time"
)

type tokenBucket struct {
	tokens int
	rate   int
	max    int
	ticker *time.Ticker
}

func (t *tokenBucket) start() {
	go func() {
		for true {
			<-t.ticker.C
			if t.tokens < t.max {
				t.tokens += t.rate
			}
			fmt.Println("令牌数量", t.tokens)
		}
	}()
}

func newTokenBucket(rate int, max int) *tokenBucket {
	return &tokenBucket{
		ticker: time.NewTicker(time.Second),
		rate:   rate,
		max:    max,
	}
}

func Ratelimit(rate int, max int) gin.HandlerFunc {
	t := newTokenBucket(rate, max)
	t.start()

	return helper.WarpHandle(func(c *gin.Context) interface{} {

		if t.tokens <= 0 {
			return responses.Forbidden("限流", nil)
		}
		t.tokens--

		return nil
	})
}
